import asyncio
import time
import random
from typing import Dict, List, Any, Callable
from functools import wraps


class FakeLLM:
    """LLM 模型，生成自然语言响应。"""

    def __init__(self, temperature: float = 0.5):
        self.temperature = temperature

    def __call__(self, prompt: str) -> str:
        """生成的响应内容，来自 GPT 模型。"""
        responses = {
            "Hello, Alice!": "Hello, Alice! Nice to meet you!",
            "Fetching data for Machine Learning.": "Here is the latest data on Machine Learning."
        }
        # 根据 prompt 返回响应
        return responses.get(prompt, "I have no idea what you are asking.")


class LangChainAgent:
    """LangChain驱动的智能体系统，支持多步骤任务链与上下文管理。"""

    def __init__(self, llm: FakeLLM):
        self.context = {}  # 上下文存储，用于任务之间的数据传递
        self.tasks = []  # 任务链容器
        self.llm = llm  # 引入 LLM 实例

    def add_task(self, func: Callable, name: str = None):
        """向任务链中添加任务，并为任务指定名称。"""
        task_name = name if name else func.__class__.__name__
        print(f"Adding task: {task_name}")
        self.tasks.append((func, task_name))

    async def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """依次执行任务链中的所有任务，并返回最终结果。"""
        data = input_data
        for task, name in self.tasks:
            print(f"Executing task: {name}")
            data = await task(data)  # 任务执行
        return data


class TemplateTask:
    """模板化任务，用于解析用户输入并生成响应。"""

    def __init__(self, template: str, llm: FakeLLM):
        self.template = template
        self.llm = llm  # 引入LLM 模型

    async def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """调用LLM 生成响应并返回结果。"""
        await asyncio.sleep(random.uniform(0.1, 0.5))
        input_value = data.get('name') or data.get('query', 'unknown')
        prompt = self.template.format(name_or_query=input_value)
        response = self.llm(prompt)  # 生成响应
        print(f"Task completed: {response}")
        return {"response": response}


class RouterChain:
    """路由链，根据条件选择任务路径。"""

    def __init__(self, routes: Dict[str, Callable]):
        self.routes = routes

    async def execute(self, condition: str, input_data: Dict[str, Any]):
        """根据条件选择并执行路径中的任务链。"""
        if condition in self.routes:
            print(f"Routing to {condition}...")
            await self.routes[condition](input_data)
        else:
            print(f"No route found for condition: {condition}")


class PerformanceMonitor:
    """性能监控器，记录并报告系统的运行情况。"""

    def __init__(self):
        self.execution_times = []

    def log_time(self, func: Callable, task_name: str = None):
        """装饰器：记录任务的执行时间，并为任务指定名称。"""
        task_name = task_name if task_name else func.__class__.__name__

        @wraps(func)
        async def wrapper(*args, **kwargs):
            start = time.time()
            result = await func(*args, **kwargs)
            elapsed = time.time() - start
            self.execution_times.append(elapsed)
            print(f"Task {task_name} executed in {elapsed:.2f}s")
            return result

        return wrapper

    def report(self):
        """生成性能报告。"""
        total = sum(self.execution_times)
        print(f"Total execution time: {total:.2f}s")
        print(f"Average execution time: {total / len(self.execution_times):.2f}s")


# 初始化 LLM 模型
llm = FakeLLM(temperature=0.5)
# 初始化智能体与性能监控器
monitor = PerformanceMonitor()
agent = LangChainAgent(llm)
# 定义任务模板并添加到任务链
greet_task = TemplateTask(template="Hello, {name_or_query}!", llm=llm)
query_task = TemplateTask(template="Fetching data for {name_or_query}.", llm=llm)
agent.add_task(monitor.log_time(greet_task, "Greet Task"))
agent.add_task(monitor.log_time(query_task, "Query ttTask"))
# 创建路由链并定义路径
router = RouterChain(routes={
    "greet": lambda data: agent.run(data),
    "query": lambda data: agent.run(data)
})


@monitor.log_time
async def main():
    """主程序入口，执行任务链并生成性能报告。"""
    print("Starting LangChain agent...")
    await router.execute("greet", {"name": "Alice"})
    await router.execute("query", {"query": "Machine Learning"})
    monitor.report()


# 执行主程序
if __name__ == "__main__":
    asyncio.run(main())
